# 演示如何利用pinocchio联合mujoco进行impedance控制
# 注意两边数据的存取与转换
# mujoco读取末端接触力以及方向，并且读取关节力矩
# 根据受力大小和方向，以及设定的目标电位，计算出末端期望速度以及方向
# 映射到关节空间


import pinocchio as pin
import mujoco
import mujoco.viewer as viewer
import numpy as np
from os.path import dirname, join
import time
from numpy.linalg import pinv

# pinocchio加载模型
modelpath = join(dirname(__file__), 'Dof6Arm.xml')
pin_model = pin.RobotWrapper.BuildFromMJCF(modelpath)

# mujoco加载模型
model = mujoco.MjModel.from_xml_path(join(dirname(__file__),'scene.xml')) 
data = mujoco.MjData(model)

key_id = model.key("home").id
mujoco.mj_resetDataKeyframe(model, data, key_id)

print(f"qpos : {data.qpos}")
print(f"type of qpos : {type(data.qpos)}")

damping: float = 1e-4
diag = damping * np.eye(6)

site_id = model.site('site_Hand').id
base_id = model.body('BigArm1').id
goal_pos = [0.5, 0.5, 0.5]
end_effector_id = model.body('Hand').id
jac = np.zeros((6, model.nv))

mocap_id = model.body("target").mocapid[0]

# pinocchio和mujoco都要更新
pin.framesForwardKinematics(pin_model.model, pin_model.data, data.qpos)
pin.updateFramePlacements(pin_model.model, pin_model.data) 
mujoco.mj_step(model, data)            # step xdot= f(x,u)
mujoco.mj_forward(model, data)

pos_sensor_names = [
        "Motor_Shoulder_pos", 
        "Motor_BigArm1_pos",
        "Motor_BigArm2_pos",
        "Motor_SmallArm1_pos",
        "Motor_SmallArm2_pos",
        "Motor_Hand_pos"
        ]
def ik(model, data):
    
    # 通过sensor获取关节角度比直接使用data.qpos更符合实际
    qpos_sensors = np.array([data.sensor(pos_sensor_name).data[0] for pos_sensor_name in pos_sensor_names])

    # 获取目标位置SE3矩阵T_t
    mocap_pos = data.mocap_pos[mocap_id]
    mocap_quat = data.mocap_quat[mocap_id]
    
    # 创建Pinocchio四元数对象（参数顺序为x,y,z,w）
    pin_quat = pin.Quaternion(mocap_quat[1], mocap_quat[2], mocap_quat[3], mocap_quat[0])
    # 转换为旋转矩阵
    mocap_rot = pin_quat.toRotationMatrix()
    T_t = pin.SE3(mocap_rot, mocap_pos)
    pin.framesForwardKinematics(pin_model.model, pin_model.data, qpos_sensors)
    pin.updateFramePlacements(pin_model.model, pin_model.data)

    # 求当前site的位姿T_b和目标位姿的相对差值矩阵T_bt
    T_bt = pin_model.data.oMf[pin_model.model.getFrameId("site_Hand")].actInv(T_t)

    # 求site_Hand的雅克比J_b
    J_b = pin.computeFrameJacobian(
        pin_model.model, 
        pin_model.data, 
        qpos_sensors, 
        pin_model.model.getFrameId("site_Hand"), 
        pin.ReferenceFrame.LOCAL_WORLD_ALIGNED  # 根据需求选择坐标系
    )

    J_l = -pin.Jlog6(T_bt.inverse())
    # 接下来会这里会利用J_error v^{*} = -e(q)来求v^{*}
    J_error = J_l.dot(J_b) # 误差Jacobian
    e_q = pin.log(T_bt).vector # 定义的误差向量
    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    v_star = pinv(J_error).dot(-e_q) # 此处是求解的关键，这里直接用pinv，没有借用damping项
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
    ctrl = v_star
    return ctrl




with viewer.launch_passive(model,data) as viewer:
    # Initialize the camera view to that of the free camera.
    mujoco.mjv_defaultFreeCamera(model, viewer.cam)
    while viewer.is_running():
        data.ctrl = ik(model,data)  # apply control
        mujoco.mj_step(model, data)            # step xdot= f(x,u)
        mujoco.mj_forward(model, data)
        viewer.sync()
        time.sleep(model.opt.timestep)